from copy import copy
import sys

class GameState:
  def expand( self ):
    raise RuntimeError
    
  def apply( self, move ):
    raise RuntimeError
    
  def eval( self ):
    raise RuntimeError
    
  def getWinner( self ):
    raise RuntimeError
    
  def getPlayer( self ):
    raise RuntimeError

class Game:
  def __init__( self, playerTypes, **moves ):
    self.playerTypes = playerTypes
    self.moves = moves
    self.moves[ 'minimax' ] = self.minimaxMove
    self.moves[ 'collect' ] = self.collectMove
    self.moves[ 'int' ] = self.intMove
    self.moves[ 'console' ] = self.consoleMove
    
  def run( self, state ):
    while not state.getWinner():
      player = state.getPlayer()
      assert player in self.playerTypes, 'unknown player %s' % player
      playerType = self.playerTypes[ player ]
      assert playerType in self.moves, 'unknown playerType %s' % playerType
      state = self.moves[ playerType ]( state )
    return state
    
  def expand( self, state ):
    return map( state.apply, state.expand() )
    
  def minimax( self, state, ply ):
    if state.getWinner() or ply <= 0:
      return state.eval()
    else:
      return max( [ -sys.maxint ] + [ -self.minimax( child, ply - 1 ) for child in self.expand( state ) ] )
    
  def minimaxMove( self, state ):
    return max( self.expand( state ), key = lambda child: self.minimax( child, 4 ) )
    
  def collect( self, state, player, ply ):
    score = state.eval()
    if state.getPlayer() == player:
      score *= -1
    scores = [ score ]
    if not state.getWinner() and ply > 0:
      for child in self.expand( state ):
        scores += self.collect( child, player, ply - 1 )
    return scores
    
  def collectMove( self, state ):
    return max( self.expand( state ), key = lambda child: sum( self.collect( child, state.getPlayer(), 2 ) ) )
    
  def stats( self, scores ):
    ave = sum( scores ) / float( len( scores ) )
    dev = sum( [ abs( score - ave ) for score in scores ] ) / float( len( scores ) )
    return ave, dev
    
  def intMove( self, state ):
    def eval( child ):
      scores = self.collect( child, state.getPlayer(), 2 )
      ave, dev = self.stats( scores )
      score = dev + ave * abs( ave )
      return score
    return max( self.expand( state ), key = eval )
    
  def consoleMove( self, state ):
    moves = state.expand()
    print 'current state:'
    print state
    print 'moves:'
    for i, move in enumerate( moves ):
      print '%d: %s' % ( i, move )
    return state.apply( moves[ input( 'enter move index: ' ) ] )

class ToeMove:
  def __init__( self, i, j, player ):
    self.i = i
    self.j = j
    self.player = player
    
  def __repr__( self ):
    return 'move player %s to %d %d' % ( self.player, self.i, self.j )
        
class Toe( GameState ):
  lines = {}

  def __init__( self, w = 3, h = 3 ):
    self.player = 'x'
    self.w = w
    self.h = h
    self.board = [ ' ' ] * self.w * self.h
    
  def __getitem__( self, i ):
    return self.board[ i[0] * self.w + i[1] ]
    
  def __setitem__( self, i, val ):
    self.board[ i[0] * self.w + i[1] ] = val
    
  def __repr__( self ):
    s = '  0 1 2\n'
    for i in range( self.h ):
      s += '%d ' % i
      for j in range( self.w ):
        s += self[j,i]
        if j < 2:
          s += '|'
        elif i < 2:
          s += '\n'
      if i < 2:
        s += '  -+-+-\n'
    return s
    
  @staticmethod
  def otherPlayer( player ):
    if player == 'x':
      return 'o'
    elif player == 'o':
      return 'x'
    else:
      raise ValueError( player )
      
  def expand( self ):
    moves = []
    for i in range( self.w ):
      for j in range( self.h ):
        if self[i,j] == ' ':
          moves.append( ToeMove( i, j, self.player ) )
    return moves
    
  def apply( self, move ):
    state = Toe()
    state.board = copy( self.board )
    state[move.i,move.j] = move.player
    state.player = Toe.otherPlayer( move.player )
    return state
    
  def getWinner( self ):
    win = min( self.w, self.h )
    if self.countLines( win, win, self.player ):
      return self.player
    elif self.countLines( win, win, Toe.otherPlayer( self.player ) ):
      return Toe.otherPlayer( self.player )
    elif all( [ val != ' ' for val in self.board ] ):
      return 'tie'
    
  def eval( self ):
    otherPlayer = self.player
    player = Toe.otherPlayer( otherPlayer )
    score = 0
    win = min( self.w, self.h )
    for p, mult in [ ( player, 1 ), ( otherPlayer, -1 ) ]:
      for i in range( 1, win + 1 ):
        score += 10 ** ( i - 1 ) * self.countLines( win, i, p )
    return score
  
  @staticmethod
  def countLines( self, length, count, player ):
    key = length, count, player, ''.join( self.board )
    if key in Toe.lines:
      return Toe.lines[key]
    numLines = 0
    otherPlayer = self.otherPlayer( player )
    for i in range( self.w ):
      for j in range( self.h ):
        for di in range( -1, 2 ):
          for dj in range( -1, 2 ):
            maxi = i + di * ( length - 1 )
            maxj = j + dj * ( length - 1 )
            if ( di or dj ) and maxi >= 0 and maxi < self.w and maxj >= 0 and maxj < self.h:
              numMatch = 0
              numOther = 0
              for k in range( length ):
                val = self[ i + di * k, j + dj * k ]
                if val == player:
                  numMatch += 1
                elif val == otherPlayer:
                  numOther += 1
              if numMatch >= count and numOther == 0:
                numLines += 1
    Toe.lines[key] = numLines
    return numLines
    
  def getPlayer( self ):
    return self.player

class ConnectFourMove:
  def __init__( self, player, i ):
    self.player = player
    self.i = i
    
  def __repr__( self ):
    return 'Add %s to column %d' % ( self.player, self.i )
    
class ConnectFour:
  def __init__( self, w = 10, h = 5 ):
    self.w = w
    self.h = h
    self.board = [ ' ' ] * self.w * self.h
    self.player = 'r'
    
  def __getitem__( self, ( i, j ) ):
    return self.board[ j * self.w + i ]
    
  def __setitem__( self, ( i, j ), val ):
    self.board[ j * self.w + i ] = val
    
  def __repr__( self ):
    s = ''
    for i in range( self.w ):
      s += '%d ' % i
    s += '\n'
    for j in range( self.h ):
      for i in range( self.w ):
        s += self[i,j]
        if i < self.w - 1:
          s += '|'
      s += '\n'
    s += '-+' * ( self.w - 1 ) + '-'
    return s
    
  def colheight( self, i ):
    j = 0
    while j < self.h and self[ i, self.h - 1 - j ] != ' ':
      j += 1
    return j
    
  def expand( self ):
    moves = []
    for i in range( self.w ):
      if self.colheight( i ) < self.h:
        moves.append( ConnectFourMove( self.player, i ) )
    return moves
    
  def apply( self, move ):
    j = self.colheight( move.i )
    assert j < self.h
    cf = ConnectFour( self.w, self.h )
    cf.board = [ val for val in self.board ]
    cf[ move.i, self.h - 1 - j ] = move.player
    cf.player = self.otherPlayer( self.player )
    return cf
    
  @staticmethod
  def otherPlayer( player ):
    if player == 'r':
      return 'b'
    else:
      return 'r'
    
  def eval( self ):
    otherPlayer = self.player
    player = ConnectFour.otherPlayer( otherPlayer )
    score = 0
    for p, m in [ ( player, 1 ), ( otherPlayer, -1 ) ]:
      for i in range( 1, 5 ):
        score += 10 ** ( i - 1 ) * Toe.countLines( self, 4, i, p )
    return score
    
  def getWinner( self ):
    otherPlayer = ConnectFour.otherPlayer( self.player )
    if Toe.countLines( self, 4, 4, self.player ):
      return self.player
    elif Toe.countLines( self, 4, 4, otherPlayer ):
      return otherPlayer
    elif all( [ val != ' ' for val in self.board ] ):
      return 'tie'
    
  def getPlayer( self ):
    return self.player
    
if __name__ == '__main__':
  final = Game( { 'r': 'console', 'b': 'minimax' } ).run( ConnectFour() )
  print final
  print 'winner:', final.getWinner()
